{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial on YOLO (V3) Training  \n",
    "\n",
    "Now let's switch the gear to YOLO_V3 Training. As a reminder, the network outputs 3 detection maps with stride 32, 16, 8 respectively （stride denotes the factor by which the image is downsampled）. Each cell of the feature map is expected to predict an object through one of it's bounding boxes containing 5+C attributes. In terms of YOLO training, we only expect the cell to predict object of which the center of ground truth object falls in its corresponding receptive field. Since each cell could predict three bounding boxes. The bounding box with higher IoU with the ground truth box is responsible for detecting the object. \n",
    "\n",
    "Now we need to understand on how to calculate the loss of boxing coordinates, objectness score, class scores. Here we explain the YOLO loss function in each of its parts. \n",
    "\n",
    "### Bounding cooridnates \n",
    "\n",
    "Loss from predicting the bounding box coordinates. Note that tx, ty, tw, th are the network outputs. The function computes a sum over each bounding box prediction of each grid cell. 1_obj is defined as 1 if an object (based on ground truth label) is expected in grid cell i and jth bounding box predictor and 0 otherwise. \n",
    "   \n",
    "### Object Confidence\n",
    "\n",
    "Loss associated with the confidence score for each bounding box predictor. x_obj is the network output for the object score. y_obj_i is equal to 1 if an object (based on ground truth label) is expected in grid cell i and jth bounding box predictor and 0 otherwise. The loss function represents a binary logistic regression cost \n",
    "\n",
    "### Classfication Loss\n",
    "\n",
    "This is essentailly also a binary logistic regression cost to represent the probabilities of the detected object belonging to a particular class (instead of softmaxing class which assume the classes are mutually exclusive, not always the case like \"Women\" and \"Person\"). y_class_i is equal to 1 for the responable bounding box predictor and the ground truth class index otherwise 0. 1_obj is involved here so classification error is not penealized when the ground truth object is not present on the bounding boxes. This is necessary to increase the model stability\n",
    "\n",
    "<img src=\"imgs/ipy_pic/eq.png\" width=\"600\" style=\"fix: right;\">\n",
    "\n",
    "The total loss is a sum of all loss terms with a hyper-parameter of weight factor that user can play with.\n",
    "\n",
    "### Training Data\n",
    "\n",
    "We use Pascal VOC dataset to demonstrate on the Loss calculation. Please refer the [YOLO Website](https://pjreddie.com/darknet/yolo/) on how to download the data and generate the label files for training. \n",
    "\n",
    "For each image, there is a paired label file. The label file shows each row for a object. The format is  [class_index], [x_center], [y_center], [width], [height]. Note that they are all normalized.\n",
    "\n",
    "### Evaluation \n",
    "\n",
    "One of the most popular approach to measure the accuracy of object detectors during training is [Precision](https://sanchom.wordpress.com/tag/average-precision) and \"[Recall](https://sanchom.wordpress.com/tag/average-precision)\". We define the \"True Positive\" as if the bounding box has IoU > 0.5 to ground truth box, the object score > 0.5 for each class. We define the \"True negative\" as all other bounding boxes which have object score > 0.5. Please be awared that this evaluation is not the same as the metrics (mAP) used by object detection competitions such as VOC and COCO. But it is necessary to apply the offical mAP to evaluate your model accuracy for either competitions or benchmarking among different approaches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 1: Data Loading and Processing\n",
    "\n",
    "Pytorch provides the 'dataloader' funcation to help users prepare the data easily. torch.utils.data.Dataset can be applied to create the dataset class. torch.utils.data.DataLoader helps batching, shuffling the data for training. Please refer to the [pytorch tutorial](https://pytorch.org/tutorials/beginner/data_loading_tutorial.html) for detailed information if you are not familiar with these functions "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run some setup code for this notebook.\n",
    "import os\n",
    "import torch\n",
    "from Darknet_VOC import Darknet\n",
    "from utils.Data_Loading import ListDataset\n",
    "from torch.utils.data import Dataset\n",
    "import time\n",
    "import copy\n",
    "import numpy as np\n",
    "import math\n",
    "import torch.nn as nn\n",
    "from utils.util import *\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first set up the training and validation path. The \"*.txt\" lists the image files as an example shown below. Here we concatenate all the 2007 and 2012 train and val files together for training (~15k images) and use 2007_test.txt(~5k images) for validation. \n",
    "\n",
    "/home/xingyu/Desktop/Project/Pascal_VOC/VOCdevkit/VOC2007/JPEGImages/000012.jpg \n",
    "\n",
    "/home/xingyu/Desktop/Project/Pascal_VOC/VOCdevkit/VOC2007/JPEGImages/000017.jpg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_path = '/home/xingyu/Desktop/Project/Pascal_VOC/train.txt'\n",
    "val_path = '/home/xingyu/Desktop/Project/Pascal_VOC/2007_test.txt'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset class is created as below. For each image path, the corresponding label file can be simplily obtained by replacing \"JPEGImages\", \".jpg\" to \"labels\" and \"txt\" \n",
    "\n",
    "OpenCV is used to load the image. The image is padded, resized and prepared to pytorch inputs as we already explained from \"Yolo_V3_Inference_Step_by_Step\" tutorial.\n",
    "\n",
    "The label has to be recalculated since we pad the image. SInce different images have different objects, in order to concatenate the label files as a batch, a zero-like label matrix with shape (50 x 5) is created. The label data is then filled in. The sample of our dataset will be a dict {'input_img': input_img, 'orig_img': pad_img, 'label': filled_labels, 'path': img_path}. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ListDataset(Dataset):\n",
    "    def __init__(self, list_path, img_size=416):\n",
    "        with open(list_path, 'r') as file:\n",
    "            self.img_files = file.readlines()\n",
    "        self.label_files = [path.replace('JPEGImages', 'labels').replace('.jpg', '.txt') for path in self.img_files]\n",
    "        self.img_shape = (img_size, img_size)\n",
    "        self.max_objects = 50\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_files)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "\n",
    "        #---------\n",
    "        #  Image\n",
    "        #---------\n",
    "\n",
    "        img_path = self.img_files[index % len(self.img_files)].rstrip()\n",
    "        img = cv2.imread(img_path)\n",
    "        h, w, _ = img.shape\n",
    "        dim_diff = np.abs(h - w)\n",
    "        # Upper (left) and lower (right) padding\n",
    "        pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2\n",
    "        # Determine padding\n",
    "        pad = ((pad1, pad2), (0, 0), (0, 0)) if h <= w else ((0, 0), (pad1, pad2), (0, 0))\n",
    "        # Add padding\n",
    "        pad_img = np.pad(img, pad, 'constant', constant_values=128)\n",
    "\n",
    "        padded_h, padded_w, _ = pad_img.shape\n",
    "\n",
    "        # Resize\n",
    "        pad_img = cv2.resize(pad_img, self.img_shape)\n",
    "        # Channels-first\n",
    "        input_img = pad_img[:, :, ::-1].transpose((2, 0, 1)).copy()\n",
    "        # As pytorch tensor\n",
    "        input_img = torch.from_numpy(input_img).float().div(255.0)\n",
    "\n",
    "        #---------\n",
    "        #  Label\n",
    "        #---------\n",
    "\n",
    "        label_path = self.label_files[index % len(self.img_files)].rstrip()\n",
    "\n",
    "        labels = None\n",
    "        if os.path.exists(label_path):\n",
    "            labels = np.loadtxt(label_path).reshape(-1, 5)\n",
    "            # Extract coordinates for unpadded + unscaled image\n",
    "            x1 = w * (labels[:, 1] - labels[:, 3]/2)\n",
    "            y1 = h * (labels[:, 2] - labels[:, 4]/2)\n",
    "            x2 = w * (labels[:, 1] + labels[:, 3]/2)\n",
    "            y2 = h * (labels[:, 2] + labels[:, 4]/2)\n",
    "            # Adjust for added padding\n",
    "            x1 += pad[1][0]\n",
    "            y1 += pad[0][0]\n",
    "            x2 += pad[1][0]\n",
    "            y2 += pad[0][0]\n",
    "            # Calculate ratios from coordinates\n",
    "            labels[:, 1] = ((x1 + x2) / 2) / padded_w\n",
    "            labels[:, 2] = ((y1 + y2) / 2) / padded_h\n",
    "            labels[:, 3] *= w / padded_w\n",
    "            labels[:, 4] *= h / padded_h\n",
    "        # Fill matrix\n",
    "        filled_labels = np.zeros((self.max_objects, 5))\n",
    "        if labels is not None:\n",
    "            filled_labels[range(len(labels))[:self.max_objects]] = labels[:self.max_objects]\n",
    "\n",
    "        filled_labels = torch.from_numpy(filled_labels)\n",
    "\n",
    "        sample = {'input_img': input_img, 'orig_img': pad_img, 'label': filled_labels, 'path': img_path}\n",
    "\n",
    "        return sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use the dataloader to load the data for both training and validaton.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 torch.Size([8, 3, 416, 416]) torch.Size([8, 416, 416, 3]) torch.Size([8, 50, 5])\n",
      "1 torch.Size([8, 3, 416, 416]) torch.Size([8, 416, 416, 3]) torch.Size([8, 50, 5])\n",
      "2 torch.Size([8, 3, 416, 416]) torch.Size([8, 416, 416, 3]) torch.Size([8, 50, 5])\n",
      "3 torch.Size([8, 3, 416, 416]) torch.Size([8, 416, 416, 3]) torch.Size([8, 50, 5])\n"
     ]
    }
   ],
   "source": [
    "batch_size = 8\n",
    "inp_dim = 416\n",
    "dataloaders = {'train': torch.utils.data.DataLoader(ListDataset(train_path, img_size=inp_dim), batch_size=batch_size, shuffle=True),\n",
    "               'val': torch.utils.data.DataLoader(ListDataset(val_path, img_size=inp_dim), batch_size=batch_size, shuffle=True)}\n",
    "\n",
    "for i_batch, sample_batched in enumerate(dataloaders[\"train\"]):\n",
    "\n",
    "    input_images_batch, Orig_images_batch, label_batch, path_batch = sample_batched['input_img'], sample_batched['orig_img'], sample_batched['label'], sample_batched['path']\n",
    "\n",
    "    print(i_batch, input_images_batch.shape, Orig_images_batch.shape, label_batch.shape)\n",
    "\n",
    "    if i_batch == 3:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 2: Loss Calculation "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to obtain the loss, we need to deduce which cell and anchor on the network output feature map is expected to predict the object. The ground truth values for the responsible bounding box will be assigned for loss calculation.\n",
    "\n",
    "The below funcation is intended to convert the ground truth label generated from dataloader to a series of matrixes which contain the ground truth info for each cell on the corresponding detection map.\n",
    "\n",
    "1. mask (representin 1_obj) is an indexing matrix with size [batch x num_anchors x grid x grid]. The responable bounding box predictor is assigned as 1, otherwise 0. \n",
    "\n",
    "2. tx, ty, tw and th has the same format as mask but containing the ground truth cooridnates instead of 1. Note that label file data here are converted to offset to centriod and log-space transfrom for width and height. \n",
    "\n",
    "3. tconf is a matrix with size [batch x num_anchors x grid x grid]. The value is 1 representing the highest probability for the responsible bounding box predictor, otherwise 0.\n",
    "\n",
    "4. tcls is a matrix with size [batch x num_anchors x grid x grid x num_class]. The value is 1 representing the highest probability for the responsible bounding box predictor and ground truth class index, otherwise 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_targets(target, anchors, grid_size, num_anchors = 3, num_classes = 20):\n",
    "\n",
    "    nB = target.size(0)\n",
    "    nA = num_anchors\n",
    "    nC = num_classes\n",
    "    nG = grid_size\n",
    "    mask = torch.zeros(nB, nA, nG, nG)\n",
    "    tx = torch.zeros(nB, nA, nG, nG)\n",
    "    ty = torch.zeros(nB, nA, nG, nG)\n",
    "    tw = torch.zeros(nB, nA, nG, nG)\n",
    "    th = torch.zeros(nB, nA, nG, nG)\n",
    "    tconf = torch.zeros(nB, nA, nG, nG)\n",
    "    tcls = torch.zeros(nB, nA, nG, nG, nC)\n",
    "\n",
    "    for b in range(nB):  # for each image\n",
    "        for t in range(target.shape[1]):  # for each object\n",
    "            if target[b, t].sum() == 0:  # if the row is empty\n",
    "                continue\n",
    "            # Convert to object label data to feature map\n",
    "            gx = target[b, t, 1] * nG\n",
    "            gy = target[b, t, 2] * nG\n",
    "            gw = target[b, t, 3] * nG\n",
    "            gh = target[b, t, 4] * nG\n",
    "            # Get grid box indices\n",
    "            gi = int(gx)\n",
    "            gj = int(gy)\n",
    "            # Get shape of gt box\n",
    "            gt_box = torch.FloatTensor(np.array([0, 0, gw, gh])).unsqueeze(0)  # 1 x 4\n",
    "            # Get shape of anchor box\n",
    "            anchor_shapes = torch.FloatTensor(\n",
    "                np.concatenate((np.zeros((len(anchors), 2)), np.array(anchors)), 1))\n",
    "            # Calculate iou between gt and anchor shapes\n",
    "            anch_ious = bbox_iou(gt_box, anchor_shapes)\n",
    "            # Find the best matching anchor box\n",
    "            best_n = np.argmax(anch_ious)\n",
    "            # Masks\n",
    "            mask[b, best_n, gj, gi] = 1\n",
    "            # Coordinates\n",
    "            tx[b, best_n, gj, gi] = gx - gi\n",
    "            ty[b, best_n, gj, gi] = gy - gj\n",
    "            # Width and height\n",
    "            tw[b, best_n, gj, gi] = math.log(gw / anchors[best_n][0] + 1e-16)\n",
    "            th[b, best_n, gj, gi] = math.log(gh / anchors[best_n][1] + 1e-16)\n",
    "            # One-hot encoding of label\n",
    "            target_label = int(target[b, t, 0])\n",
    "            tcls[b, best_n, gj, gi, target_label] = 1\n",
    "            tconf[b, best_n, gj, gi] = 1\n",
    "\n",
    "    return mask, tx, ty, tw, th, tconf, tcls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss funcation is applied to calculate loss between one of the dectection map and ground truth label converted using \"build target\" funcation. The network output including center offset (x, y), dimension log-space transform (w, h), object score - pred_conf, class score - pred_cls are reshaped as the same maxtixes as ground truth feature map. Each loss item and overall loss are then calculated based on the equation shown before.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Loss(input, target, anchors, inp_dim, num_anchors = 3, num_classes = 20):\n",
    "\n",
    "    nA = num_anchors  # number of anchors\n",
    "    nB = input.size(0)  # number of batches\n",
    "    nG = input.size(2)  # number of grid size\n",
    "    nC = num_classes\n",
    "    stride = inp_dim / nG\n",
    "\n",
    "    # Tensors for cuda support\n",
    "    FloatTensor = torch.cuda.FloatTensor if input.is_cuda else torch.FloatTensor\n",
    "    ByteTensor = torch.cuda.ByteTensor if input.is_cuda else torch.ByteTensor\n",
    "\n",
    "    prediction = input.view(nB, nA, 5 + nC, nG, nG).permute(0, 1, 3, 4, 2).contiguous()  # reshape the output data\n",
    "\n",
    "    # Get outputs\n",
    "    x = torch.sigmoid(prediction[..., 0])  # Center x\n",
    "    y = torch.sigmoid(prediction[..., 1])  # Center y\n",
    "    w = prediction[..., 2]  # Width\n",
    "    h = prediction[..., 3]  # Height\n",
    "    pred_conf = torch.sigmoid(prediction[..., 4])  # Conf\n",
    "    pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred\n",
    "\n",
    "    # Calculate offsets for each grid\n",
    "    grid_x = torch.arange(nG).repeat(nG, 1).view([1, 1, nG, nG]).type(FloatTensor)\n",
    "    grid_y = torch.arange(nG).repeat(nG, 1).t().view([1, 1, nG, nG]).type(FloatTensor)\n",
    "    scaled_anchors = FloatTensor([(a_w / stride, a_h / stride) for a_w, a_h in anchors])\n",
    "    anchor_w = scaled_anchors[:, 0:1].view((1, nA, 1, 1))\n",
    "    anchor_h = scaled_anchors[:, 1:2].view((1, nA, 1, 1))\n",
    "\n",
    "    # Add offset and scale with anchors\n",
    "    pred_boxes = FloatTensor(prediction[..., :4].shape)\n",
    "    pred_boxes[..., 0] = x.data + grid_x\n",
    "    pred_boxes[..., 1] = y.data + grid_y\n",
    "    pred_boxes[..., 2] = torch.exp(w.data) * anchor_w\n",
    "    pred_boxes[..., 3] = torch.exp(h.data) * anchor_h\n",
    "\n",
    "    mask, tx, ty, tw, th, tconf, tcls = build_targets(\n",
    "        target=target.cpu().data,\n",
    "        anchors=scaled_anchors.cpu().data,\n",
    "        grid_size=nG,\n",
    "        num_anchors=nA,\n",
    "        num_classes=num_classes)\n",
    "\n",
    "    # Handle target variables\n",
    "    tx, ty = tx.type(FloatTensor), ty.type(FloatTensor)\n",
    "    tw, th = tw.type(FloatTensor), th.type(FloatTensor)\n",
    "    tconf, tcls = tconf.type(FloatTensor), tcls.type(FloatTensor)\n",
    "    mask = mask.type(ByteTensor)\n",
    "\n",
    "    mse_loss = nn.MSELoss(reduction='sum')  # Coordinate loss\n",
    "    bce_loss = nn.BCELoss(reduction='sum')  # Confidence loss\n",
    "    loss_x = mse_loss(x[mask], tx[mask])\n",
    "    loss_y = mse_loss(y[mask], ty[mask])\n",
    "    loss_w = mse_loss(w[mask], tw[mask])\n",
    "    loss_h = mse_loss(h[mask], th[mask])\n",
    "    loss_conf = bce_loss(pred_conf, tconf)\n",
    "    loss_cls = bce_loss(pred_cls[mask], tcls[mask])\n",
    "    loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls\n",
    "\n",
    "    return (loss, loss_x, loss_y, loss_w, loss_h, loss_conf, loss_cls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is a loss calculation test using the pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'total_loss': tensor(81.8016, device='cuda:0', grad_fn=<AddBackward0>),\n",
       " 'x': tensor(2.2884, device='cuda:0', grad_fn=<AddBackward0>),\n",
       " 'y': tensor(2.1074, device='cuda:0', grad_fn=<AddBackward0>),\n",
       " 'w': tensor(1.6042, device='cuda:0', grad_fn=<AddBackward0>),\n",
       " 'h': tensor(1.2552, device='cuda:0', grad_fn=<AddBackward0>),\n",
       " 'conf': tensor(69.0602, device='cuda:0', grad_fn=<AddBackward0>),\n",
       " 'cls': tensor(5.4863, device='cuda:0', grad_fn=<AddBackward0>)}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = Darknet()\n",
    "model.load_state_dict(torch.load('weights/Dartnet_VOC_Weights'))\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "CUDA = torch.cuda.is_available()\n",
    "model = model.to(device)\n",
    "\n",
    "inputs = input_images_batch.to(device)\n",
    "labels = label_batch.to(device)\n",
    "Final_pre, *output = model(inputs, CUDA)\n",
    "\n",
    "anchors = ([(10, 13), (16, 30), (33, 23)], [(30, 61), (62, 45), (59, 119)], [(116, 90), (156, 198), (373, 326)])\n",
    "loss_item = {\"total_loss\": 0, \"x\": 0, \"y\": 0, \"w\": 0, \"h\": 0, \"conf\": 0, \"cls\": 0}\n",
    "\n",
    "for i in range(len(output)):\n",
    "    losses = Loss(output[i], labels.float(), anchors[i], inp_dim=inp_dim, num_anchors = 3, num_classes = 20)\n",
    "    for i, name in enumerate(loss_item):\n",
    "        loss_item[name] += losses[i]\n",
    "loss_item"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 3: Evaluation \n",
    "\n",
    "The evaulation is crucial during training to pick up the best performance weights. The \"eval\" function is built to evaluate the model performance. The funcation accepts the model prediction after object confidence and NMS filters (output from function \"write_results\") and ground truth label data for dataloader, return the batch precision, recall and F1 score.\n",
    "\n",
    "The number of batch ground truth can be directly derived from ground truth label data\n",
    "\n",
    "The total number of proposals (true positive + negative positive) is calculated for all bounding boxes with object score > 0.5\n",
    "\n",
    "The number of true positive is calculated if IoU against ground truth bounding box is higher than 0.5 for each class. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval(output, labels, img_width, img_height):\n",
    "\n",
    "    nProposals = int((output[:, 5] > 0.5).sum().item())\n",
    "    nGT = 0\n",
    "    nCorrect = 0\n",
    "    for b in range(labels.shape[0]):  # for each image\n",
    "        prediction = output[output[:,0] == b]  # filter out the predictions of corresponding image\n",
    "        for t in range(labels.shape[1]):  # for each object\n",
    "            if labels[b, t].sum() == 0:  # if the row is empty\n",
    "                continue\n",
    "            nGT += 1\n",
    "            gt_label = convert_label(labels[b, t].unsqueeze(0), img_width, img_height)\n",
    "            gt_box = gt_label[:, 1:5]\n",
    "            for i in range(prediction.shape[0]):\n",
    "                pred_box = prediction[i, 1:5].unsqueeze(0)\n",
    "                iou = bbox_iou(pred_box, gt_box)\n",
    "                pred_label = prediction[i, -1]\n",
    "                target_label = gt_label[0, 0]\n",
    "                if iou > 0.5 and pred_label == target_label:\n",
    "                    nCorrect += 1\n",
    "    recall = float(nCorrect / nGT) if nGT else 1\n",
    "    precision = float(nCorrect / nProposals) if nProposals else 0\n",
    "    F1_score = 2 * recall * precision / (recall + precision + 1e-16)\n",
    "\n",
    "    return F1_score, precision, recall"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_label(image_anno, img_width, img_height):\n",
    "\n",
    "    \"\"\"\n",
    "    Function: convert image annotation : center x, center y, w, h (normalized) to x1, y1, x2, y2 for corresponding img\n",
    "    \"\"\"\n",
    "    x_center = image_anno[:, 1]\n",
    "    y_center = image_anno[:, 2]\n",
    "    width = image_anno[:, 3]\n",
    "    height = image_anno[:, 4]\n",
    "\n",
    "    output = torch.zeros_like(image_anno)\n",
    "    output[:,0] = image_anno[:,0]\n",
    "    output[:, 1], output[:, 3] = x_center - width / 2, x_center + width / 2\n",
    "    output[:, 2], output[:, 4] = y_center - height / 2, y_center + height / 2\n",
    "\n",
    "    output[:, [1, 3]] *= img_width\n",
    "    output[:, [2, 4]] *= img_height\n",
    "\n",
    "    return output.type(torch.FloatTensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's use the pretrained model to test the results "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Recall: 0.8750 Precision: 1.0000 F1 Score: 0.9333\n"
     ]
    }
   ],
   "source": [
    "Final_pre = write_results(Final_pre, confidence=0.5, num_classes=20, nms_conf=0.4)\n",
    "\n",
    "if isinstance(Final_pre, int) == False:\n",
    "    F1_score, precision, recall = eval(Final_pre.cpu(), labels, img_width=inp_dim, img_height=inp_dim)\n",
    "else:\n",
    "    F1_score, precision, recall = 0, 0, 0\n",
    "\n",
    "print('Recall: {:.4f} Precision: {:.4f} F1 Score: {:.4f}'.format(recall, precision, F1_score))                                                                                    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Section 4:  Training \n",
    "\n",
    "So far, we already prepared all the necessay toolkits for training. We will train the network using transfer learning. The pretrained Yolo_V3 weights can be downloaded from [YOLO Website](https://pjreddie.com/darknet/yolo/). we can use the this weights either as an initialization or a fixed feature extractor. As learned from Standard cs231n course, since we have very similar dataset but not quite a lot of data, it is approporiate to freeze the weights for all of the network expect the last three layers for each detection. However, if you have other types of datasets, you can switch your strategy based on the task of interests. \n",
    "\n",
    "In this project we already generate the initialization weights, you can just download it from this [link](https://pan.baidu.com/s/1-O-jD0uU3OM6yNaUSLjAhw#list/path=%2F).\n",
    "\n",
    "We use the Adam optimization algorithm as the default choice, SGD+Momentum with learning rate decay might be worth to try if you are looking for approaches to drive more accuracy, but it requres more tuning. \n",
    "\n",
    "F1 score is eventually applied to pick up the best weights. The results for epoch 0 are presented below. The overall training takes 5 hours using Nvidia Quadro P4000 for 24 epochs "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0/0\n",
      "----------\n",
      "train Loss: 62.8359 Recall: 0.4240 Precision: 0.7345 F1 Score: 0.5238\n",
      "train xy: 0.8823 wh: 1.5822 conf: 54.4661 class: 5.9054\n",
      "val Loss: 26.4952 Recall: 0.5652 Precision: 0.7699 F1 Score: 0.6418\n",
      "val xy: 0.7917 wh: 0.9582 conf: 21.5911 class: 3.1542\n",
      "Training complete in 14m 27s\n",
      "Best F1 score: 0.641817\n"
     ]
    }
   ],
   "source": [
    "model = Darknet()\n",
    "model.load_state_dict(torch.load('weights/Dartnet_VOC_weights_ini'))\n",
    "\n",
    "#  This section is to freeze all the network except the output three layers\n",
    "for name, param in model.named_parameters():\n",
    "    param.requires_grad = False\n",
    "    if int(name.split('.')[1]) in (79, 80, 81, 91, 92, 93, 103, 104, 105):\n",
    "        param.requires_grad = True\n",
    "\n",
    "model = model.to(device)\n",
    "optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()))\n",
    "\n",
    "since = time.time()\n",
    "\n",
    "best_model_wts = copy.deepcopy(model.state_dict())\n",
    "best_F1_score = 0.0\n",
    "\n",
    "num_epochs = 1\n",
    "for epoch in range(num_epochs):\n",
    "    print('Epoch {}/{}'.format(epoch, num_epochs-1))\n",
    "    print('-' * 10)\n",
    "\n",
    "    # Each epoch has a training and validation phase\n",
    "    for phase in ['train', 'val']:\n",
    "        if phase == 'train':\n",
    "            model.train() # set model to training mode\n",
    "        else:\n",
    "            model.eval() # set model to evaluate mode\n",
    "\n",
    "        running_loss, running_xy_loss, running_wh_loss, running_conf_loss, running_cls_loss = 0.0, 0.0, 0.0, 0.0, 0.0\n",
    "        running_recall, running_precision, running_F1_score = 0.0, 0.0, 0.0\n",
    "\n",
    "        # iterate over data\n",
    "        for i_batch, sample_batched in enumerate(dataloaders[phase]):\n",
    "            inputs, labels = sample_batched['input_img'], sample_batched['label']\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # zero the parameter gradients\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # forward\n",
    "            # track history if only in train\n",
    "            with torch.set_grad_enabled(phase == 'train'):\n",
    "\n",
    "                Final_pre, *output = model(inputs, CUDA)\n",
    "                Final_pre = write_results(Final_pre, confidence=0.5, num_classes=20, nms_conf=0.4)\n",
    "\n",
    "                anchors = (\n",
    "                [(10, 13), (16, 30), (33, 23)], [(30, 61), (62, 45), (59, 119)], [(116, 90), (156, 198), (373, 326)])\n",
    "\n",
    "                loss_item = {\"total_loss\": 0, \"x\": 0, \"y\": 0, \"w\": 0, \"h\": 0, \"conf\": 0, \"cls\": 0}\n",
    "\n",
    "                for i in range(len(output)):\n",
    "                    losses = Loss(output[i], labels.float(), anchors[i], inp_dim=inp_dim, num_anchors = 3, num_classes = 20)\n",
    "                    for i, name in enumerate(loss_item):\n",
    "                        loss_item[name] += losses[i]\n",
    "\n",
    "                if isinstance(Final_pre, int) == False:\n",
    "                    F1_score, precision, recall = eval(Final_pre.cpu(), labels, img_width=inp_dim, img_height=inp_dim)\n",
    "                else:\n",
    "                    F1_score, precision, recall = 0, 0, 0\n",
    "\n",
    "                loss = loss_item['total_loss']\n",
    "                xy_loss = loss_item['x']+loss_item['y']\n",
    "                wh_loss = loss_item['w']+loss_item['h']\n",
    "                conf_loss = loss_item['conf']\n",
    "                cls_loss = loss_item['cls']\n",
    "\n",
    "                # backward + optimize only if in training phase\n",
    "                if phase == 'train':\n",
    "                    loss.backward()\n",
    "                    optimizer.step()\n",
    "\n",
    "                # statistics\n",
    "                running_loss += loss.item()\n",
    "                running_xy_loss += xy_loss.item()\n",
    "                running_wh_loss += wh_loss.item()\n",
    "                running_conf_loss += conf_loss.item()\n",
    "                running_cls_loss += cls_loss.item()\n",
    "                running_recall += recall\n",
    "                running_precision += precision\n",
    "                running_F1_score += F1_score\n",
    "\n",
    "        epoch_loss = running_loss / ((i_batch+1)*batch_size)\n",
    "        epoch_xy_loss = running_xy_loss / ((i_batch+1)*batch_size)\n",
    "        epoch_wh_loss = running_wh_loss / ((i_batch + 1) * batch_size)\n",
    "        epoch_conf_loss = running_conf_loss / ((i_batch + 1) * batch_size)\n",
    "        epoch_cls_loss = running_cls_loss / ((i_batch + 1) * batch_size)\n",
    "\n",
    "        epoch_recall = running_recall / (i_batch+1)\n",
    "        epoch_precision = running_precision / (i_batch+1)\n",
    "        epoch_F1_score = running_F1_score / (i_batch+1)\n",
    "\n",
    "        print(\n",
    "            '{} Loss: {:.4f} Recall: {:.4f} Precision: {:.4f} F1 Score: {:.4f}'.format(phase, epoch_loss, epoch_recall,\n",
    "                                                                                       epoch_precision, epoch_F1_score))\n",
    "        print(\n",
    "            '{} xy: {:.4f} wh: {:.4f} conf: {:.4f} class: {:.4f}'.format(phase, epoch_xy_loss, epoch_wh_loss,\n",
    "                                                                                                        epoch_conf_loss,\n",
    "                                                                                                        epoch_cls_loss))\n",
    "\n",
    "        # deep copy the model\n",
    "        if phase == 'val' and epoch_F1_score > best_F1_score:\n",
    "            best_F1_score = epoch_F1_score\n",
    "            best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "time_elapsed = time.time() - since\n",
    "print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))\n",
    "print('Best F1 score: {:4f}'.format(best_F1_score))\n",
    "\n",
    "model.load_state_dict(best_model_wts)\n",
    "torch.save(model.state_dict(), 'Dartnet_VOC_Weights')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
